#!/usr/bin/env python
# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: neo

import sys
import os
import subprocess


def get_path_component(a_path):
  res = []
  while a_path != '':
    a_path, tail = os.path.split(a_path)
    res.insert(0, tail)
  return res


def get_env():
  env = {}
  env.update(os.environ)
  if 'TZ' not in env:
    env['TZ'] = 'UTC'
  extra_python_paths = [
    '',
    '../',
    '../bazel-bin',
    '../external/xunkemgmt/xunkemgmt_client/python',
  ]
  if sys.version_info >= (3, 10, 0):
    extra_python_paths.append('../../coin_binary/coin2-bazel-bin-3.10')
    extra_python_paths.append('../../coin_binary/coin2-bazel-bin-3.10/external/apiset')
    extra_python_paths.append('../../coin_binary/motion_lib/pyorio-0.0.1-py3.10-linux-x86_64.egg')
    extra_python_paths.append('../../coin_binary/motion_lib/pypapyrus-0.0.1-py3.10-linux-x86_64.egg')
  else:
    extra_python_paths.append('../../coin_binary/coin2-bazel-bin')
    extra_python_paths.append('../../coin_binary/coin2-bazel-bin/external/apiset')
    extra_python_paths.append('../../coin_binary/motion_lib/pyorio-0.0.1-py3.8-linux-x86_64.egg')
  extra_python_paths = [os.path.normpath(os.path.join(os.getcwd(), path))
                        for path in extra_python_paths]
  python_path = ':'.join(extra_python_paths)
  if 'PYTHONPATH' in env:
    env['PYTHONPATH'] += ':' + python_path
  else:
    env['PYTHONPATH'] = python_path
  return env
  '''TODO(jaewon): Add custom env vars.
  For example:

    if 'LD_LIBRARY_PATH' in env:
      env['LD_LIBRARY_PATH'] += ':' + lib_path
    else:
      env['LD_LIBRARY_PATH'] = lib_path
  '''
  return env


def main(argv):
  if '--pdb' in argv:
    sys.stderr.write('[pyrunner] setup pdb.post_morten')
    pdb = True
    argv.remove('--pdb')
  else:
    pdb = False

  if len(argv) < 2:
    sys.stderr.write("No python file is given.\n")
    sys.stderr.write("Usage: ./pyrunner <python file>\n")
    return 1

  py_path = os.path.normpath(argv[1])
  assert os.path.isfile(py_path), py_path
  assert not os.path.isabs(py_path), py_path

  a_path, ext = os.path.splitext(py_path)
  comp = get_path_component(a_path)
  assert ext == '.py', ext
  assert len(comp) >= 1, len(comp)
  module_path = '.'.join(comp)

  if pdb:
    cmd = f"import pdb; import runpy; module_path='{module_path}'; exec('try: runpy.run_module(module_path, run_name=\\'__main__\\')\\nexcept: pdb.post_mortem()')"
    args = ['python', '-c', cmd] + argv[2:]
  else:
    args = ['python', '-m', module_path] + argv[2:]

  os.execvpe('python', args, get_env())


if __name__ == '__main__':
  sys.exit(main(sys.argv))
